import os

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from tqdm import tqdm
from functools import partial
from tools import par_run, par_wrapper
import threading
import copy
import multiprocessing
import pickle
import scipy.linalg as LA

# TODO：1）绘制单仿真角度，四种方法的多次仿真效果，绘制点图横轴为方法，纵轴为角度
# TODO：2）完成奇异值分解降噪，对比降噪前后DOA估计效果（评估一下是否可视化噪声曲线和降噪曲线）
# TODO：3）执行不同快拍数的仿真（服务器）
# TODO：4）平均声强法估计效果（可以对比阵列方法）阵列方法好像只有声强？矢量测量量有哪些


class DOA():
    def __init__(self):
        # 参数设置
        self.num_signals = 2
        self.num_sensors = 4
        self.sensor_distance = 0.5  # 传感器间距
        self.wavelength = 1  # 波长
        self.SNR_dB = 20  # 信噪比
        self.spectrum = None
        self.theta_resolution = 0.2  # MUSIC使用
        self.evaluate_data = {}
        self.workers = 60
        self.num_tests = 1000
        self.len_s_sample = 0
        self.default_len_s = 1000
        # 在平均声强法生成信号时考虑了采样频率和信号频率不同
        self.fs = 8000  # 采样频率
        self.f = 1000  # 信号频率

    # 生成模拟水听器阵列信号
    def generate_signal(self, SNR=None, print_angle=True, len_signal=1000, fixed_angle=None):
        '''
        return:
        array_signal:(num_sensors, num_samples) 输出仿真信号
        '''
        self.SNR_dB = SNR if SNR is not None else self.SNR_dB
        noise_power = 10 ** (-self.SNR_dB / 10)
        signal_angles = np.random.uniform(0, np.pi, self.num_signals)
        if fixed_angle is not None:
            signal_angles = np.array(fixed_angle)
        if print_angle:
            for i, angle in enumerate(sorted(signal_angles)):
                print(f"信号{i}的角度真值为：{np.rad2deg(angle)}")

        # 生成模拟水听器阵列信号
        array_signal = np.zeros((self.num_sensors, len_signal), dtype=complex)
        d_over_lambda = self.sensor_distance / self.wavelength

        for signal_idx, angle_truth in enumerate(signal_angles):
            for i in range(len_signal):
                source_signal = np.exp(1j * np.random.uniform(0, 2 * np.pi))  # 随机相位的声源信号
                noise = np.random.normal(0, np.sqrt(noise_power / 2), size=(self.num_sensors,)) + \
                        1j * np.random.normal(0, np.sqrt(noise_power / 2), size=(self.num_sensors,))
                array_signal[:, i] += np.exp(
                    1j * 2 * np.pi * np.arange(self.num_sensors) * np.cos(
                        angle_truth) * d_over_lambda) * source_signal + noise
        # array_signal = np.real(array_signal)
        return array_signal, np.sort(signal_angles)

    def read_signal(self, data_dir, angle_file, signal_file):
        angle_data = np.loadtxt(os.path.join(data_dir, "阵列罗盘航向姿态角数据", angle_file))
        signal_data = np.loadtxt(os.path.join(data_dir, "换能器信号数据", signal_file))

        # 计算角度平均值
        angle = np.mean(angle_data)

        # 整合信号数据
        num_sensors = 3  # 三元矢量水听器
        num_samples = signal_data.shape[0]
        array_signal = signal_data.reshape(num_samples, num_sensors, 3)  # (num_samples, num_sensors, 3)

        # 按照每个水听器构建复数信号
        complex_signals = []
        for i in range(num_sensors):
            complex_signal = array_signal[:, i, 0] + 1j * array_signal[:, i, 1]  # x + 1j * y
            complex_signals.append(complex_signal)

        array_signal = np.array(complex_signals)  # (num_sensors, num_samples)
        return array_signal, angle

    def generate_signal_mean_intensity(self, theta, len_s, SNR):
        """
        生成指定方位的信号
        """
        # print(f"生成的信号方向为：{np.rad2deg(theta)}")
        t = np.arange(len_s) / self.fs  # 采样时刻序列
        pressure = np.sin(2 * np.pi * self.f * t)
        velocity_x = np.cos(theta) * pressure  # x轴振速与声压同相
        velocity_y = np.sin(theta) * pressure  # y轴振速与声压同相
        noise_power = 10 ** (-SNR / 10)
        noise_p = np.random.randn(len_s) * np.sqrt(noise_power)
        noise_nx = np.random.randn(len_s) * np.sqrt(noise_power / 2)
        noise_ny = np.random.randn(len_s) * np.sqrt(noise_power / 2)

        return pressure + noise_p, velocity_x + noise_nx, velocity_y + noise_ny

    def mean_intensity_2(self, pressure, velocity_x, velocity_y):
        """
        平均声强法估计DOA
        """
        # 计算x轴和y轴的声强
        intensity_x = pressure * velocity_x
        intensity_y = pressure * velocity_y

        # 计算平均声强
        mean_intensity_x = np.mean(intensity_x)
        mean_intensity_y = np.mean(intensity_y)

        # 估计DOA
        theta_est = np.arctan2(mean_intensity_y, mean_intensity_x)

        return theta_est

    def mean_intensity(self, signal):
        assert isinstance(signal, tuple)
        pressure, velocity_x, velocity_y = signal
        # 计算x轴和y轴的声强
        intensity_x = pressure * velocity_x
        intensity_y = pressure * velocity_y

        # 计算平均声强
        mean_intensity_x = np.mean(intensity_x)
        mean_intensity_y = np.mean(intensity_y)

        # 估计DOA
        theta_est = np.arctan2(mean_intensity_y, mean_intensity_x)

        return theta_est



    def esprit(self, array_signal):
        # 计算协方差矩阵
        R = np.dot(array_signal, array_signal.conj().T) / array_signal.shape[1]
        # 特征值分解并取得信号子空间
        D, U = np.linalg.eigh(R)  # D为特征值， U为特征向量
        Us = U[:, -self.num_signals:]  # 信号子空间
        # 角度估计
        Ux = Us[:-1, :]  # Us为M*M矩阵，前M-1构成Ux  Ux(sensors-1, signal)
        Uy = Us[1:, :]
        # 最小二乘返回(x, residuals, rank, s)，其中 x 是解向量
        Psi = np.linalg.lstsq(Ux, Uy, rcond=None)[0]  # 最小二乘解函数Ux * Phi = Uy
        T, Phi = np.linalg.eig(Psi)  # T是特征值，Phi是特征向量 phi=(signal,signal)
        # Theta = np.rad2deg(np.arccos(np.angle(np.diag(Phi))))
        # Theta = np.arccos(np.angle(np.diag(Phi)) * self.wavelength / (2 * np.pi * self.sensor_distance))  # 估计角度
        Theta = np.arccos(np.angle(T) * self.wavelength / (2 * np.pi * self.sensor_distance))  # 估计角度
        Theta = np.sort(Theta)
        return Theta

    def TLS_ESPRIT(self, array_signal):
        #计算协方差矩阵
        R = np.dot(array_signal, array_signal.conj().T) / array_signal.shape[1]
        # 特征值分解
        D, EV = LA.eigh(R)
        index = np.argsort(D)
        EN = EV[:, index][:, -self.num_signals:]

        Exy = np.concatenate((EN[:-1, :], EN[1:, :]), axis=1)
        E_xy = Exy.conj().T @ Exy
        D_xy, EV_xy = LA.eig(E_xy)
        index1 = np.argsort(D_xy)
        EV_xy = EV_xy[:, index1]
        E12 = EV_xy[:self.num_signals, :self.num_signals]
        E22 = EV_xy[self.num_signals:, :self.num_signals]
        Psi = -E12 @ LA.pinv(E22)

        D_si, V_si = LA.eig(Psi)
        ephi = np.angle(D_si)
        angles = np.arccos(ephi / (2 * np.pi) / (self.sensor_distance / self.wavelength))
        return np.sort(angles)

    # 定义MUSIC算法函数
    def music(self, signal_array):
        # 获取信号个数和传感器数量
        num_sensors, num_samples = signal_array.shape
        # 计算空间参数
        omega = 2 * np.pi * self.sensor_distance / self.wavelength * np.cos(
            np.deg2rad(np.arange(0, 180, self.theta_resolution)))
        # 计算信号协方差矩阵
        covariance_matrix = np.cov(signal_array)
        # 对协方差矩阵进行特征值分解
        eigenvalues, eigenvectors = np.linalg.eigh(covariance_matrix)  # 特征值从小到大排列

        # 选取前K个特征值和对应的特征向量
        dim = num_sensors - self.num_signals
        noisy_subspace = eigenvectors[:, :dim]

        # 初始化谱函数
        self.spectrum = np.zeros(len(omega))  # len = 180

        # 计算谱函数
        for i in range(len(omega)):  # 遍历每一个谱角度
            a_theta = np.exp(1j * omega[i] * np.arange(num_sensors)).reshape(4, 1)
            A = a_theta.conj().T
            B = noisy_subspace
            C = noisy_subspace.conj().T
            D = a_theta
            self.spectrum[i] = 1 / np.abs(np.dot(np.dot(np.dot(A, B), C), D))

        # 寻找波达方向的估计值
        # peak_index = np.argsort(spectrum)[-num_signals:]
        # peak_index, properties = find_peaks(self.spectrum, distance=3, prominence=1)
        peak_index, properties = find_peaks(self.spectrum, distance=3, prominence=0)

        # estimated_theta = np.rad2deg(np.arccos(omega[peak_index] * wavelength / (2 * np.pi * sensor_distance)))
        peak_heights = properties["prominences"]
        sorted_indices = np.argsort(peak_heights)[::-1]  # 获得从大到小的索引
        # 取前num_signals个峰值
        final_index = sorted_indices[:min(len(sorted_indices), self.num_signals)]
        top_peak_indices = peak_index[final_index]
        self.peak_index = top_peak_indices
        angle = np.sort(np.deg2rad(top_peak_indices*self.theta_resolution))
        return angle

    # MUSIC算法绘图
    def draw_spectrum_result(self, spectrum, estimated_theta):
        # 绘制声源定向结果
        plt.figure()
        plt.plot(np.arange(0, 180, self.theta_resolution), spectrum)
        for i, theta_rad in enumerate(estimated_theta):
            angle = np.rad2deg(theta_rad)
            print(f"估计的角度{i}为：{angle:.3f}")
            plt.axvline(x=angle, color='r', linestyle='--', linewidth=1)
            plt.text(angle + 12, -np.max(spectrum) / 25, f'x={angle:.3f}', color='r', fontsize=12, ha='center')
            spectrum_index = round(angle/self.theta_resolution)
            plt.plot(angle, spectrum[spectrum_index], "o", label="peaks")
        for theta in self.truth_theta:
            angle = np.rad2deg(theta)
            plt.axvline(x=angle, color='g', linestyle='--', linewidth=1)


        plt.title(f'MUSIC Algorithm for Source Localization  ({self.SNR_dB}dB)')
        plt.xlabel('Theta (degrees)')
        plt.ylabel('Spectrum')
        plt.grid(True)
        plt.show()

    # root-MUSIC算法函数
    def root_music(self, array_signal):
        # root-MUSIC 算法
        M, T = array_signal.shape  # (num_sensors, num_samples)
        K = self.num_signals  # 信源数

        # 计算协方差矩阵self.
        R = np.cov(array_signal)
        # 特征值分解并取得噪声子空间
        D, U = np.linalg.eigh(R)  # D为特征值， U为特征向量
        Un = U[:, :K]
        Gn = Un @ Un.conj().T
        # 提取多项式系数并对多项式求根
        coe = np.zeros(2 * M - 1, dtype=complex)
        for i in range(-(M - 1), M):  # M -> sensors
            coe[-i + M - 1] = np.sum(np.diag(Gn, i))
        r = np.roots(coe)  # 利用roots函数求多项式的根
        r = r[np.abs(r) < 1]  # 找出在单位圆里的根
        distances = np.abs(np.abs(r) - 1)
        sorted_indices = np.argsort(distances)
        sorted_r = r[sorted_indices][:K]
        Theta = np.arccos(np.angle(sorted_r) * self.wavelength / (2 * np.pi * self.sensor_distance))  # 计算信号到达方向角
        return np.sort(Theta)

    def apply_MUSIC(self, fixed_angle):

        # 生成模拟水听器阵列信号
        array_signal, truth_theta = self.generate_signal(fixed_angle=fixed_angle, print_angle=True)

        self.truth_theta = truth_theta
        # 运行MUSIC算法
        estimated_theta = self.music(array_signal)

        # 绘制角度谱结果
        self.draw_spectrum_result(self.spectrum, estimated_theta)

    def apply_doa_methods(self, methods):
        # 添加绘图可视化，同一个真值角度用四个方法仿真多次，横轴为每个方法，纵轴为角度，每个方法的位置都绘制真值点
        array_signal, truth_angles = self.generate_signal()
        for doa_method in methods:
            print(f"========={doa_method.__name__}=========")
            estimated_theta = doa_method(array_signal)
            for i, theta in enumerate(estimated_theta):
                angle = np.rad2deg(theta)
                print(f"估计的角度{i}为：{angle:.3f}")

    def evaluation_SNR(self, doa_method, snr_range=(-10, 30), snr_sample=100):
        SNR_list = np.linspace(snr_range[0], snr_range[1], snr_sample)  # 等间隔采样
        wrong_esti_signal_num = 0
        rmse_list = []
        for index, snr in enumerate(tqdm(SNR_list, desc='Processing')):
            # for index, snr in enumerate(SNR_list):  # 从小到大遍历所有SNR采样
            total_errors = 0
            count_test_signal = 0
            for _ in range(self.num_tests):
                signal, truth_angels = self.generate_signal(snr, print_angle=False)
                estimated_theta = doa_method(signal)
                estimated_theta = np.sort(estimated_theta)
                # 搜索角度和真值不对应暂时先不考虑
                if len(estimated_theta) != len(truth_angels):
                    wrong_esti_signal_num += 1
                    continue
                for i, theta in enumerate(estimated_theta):
                    count_test_signal += 1
                    total_errors += (theta - truth_angels[i]) ** 2
            if count_test_signal > 0:
                rmse_list.append(total_errors / count_test_signal)
            else:
                SNR_list = np.delete(SNR_list, index)
                print(f"test{index},count zero!  SNR={snr:.4f}")

        print(f"error signals number:{wrong_esti_signal_num}！")
        return SNR_list, rmse_list

    def plot_from_file(self, num_tests, snr_sample, method_list=("root_music", "esprit"), task="snr"):
        plt.figure()
        for method in method_list:
            pkl_path = f"{task}/{method}_{num_tests}test_{snr_sample}sample.pkl"
            with open(pkl_path, "rb") as f:
                loaded_data = pickle.load(f)
            SNR_list = loaded_data[0, :]
            rmse_list = loaded_data[1, :]
            plt.plot(SNR_list, rmse_list, label=f'{method}')
        plt.title(f'RMSE SNR Evaluation {snr_sample} samples {num_tests} tests')
        plt.xlabel('SNR/dB')
        plt.ylabel('RMSE')
        plt.legend()
        plt.grid(True)
        plt.show()

    def apply_mean_intensity_doa(self, num_trials, len_s=1024):
        '''
        应用平均声强法进行仿真实验，单独实现的仿真信号生成
        Returns:
        解算角度、平均解算误差
        '''
        errors = []
        angles = []
        for _ in range(num_trials):
            # 生成随机方位角
            theta_true = np.random.uniform(-np.pi / 2, np.pi / 2)
            SNR = 10
            # 生成信号
            p, v_x, v_y = self.generate_signal_mean_intensity(theta_true, len_s, SNR)
            # signal2 = generate_signal(theta_true, 1024)

            # 估计DOA
            theta_est = self.mean_intensity_2(p, v_x, v_y)
            print(f"估计的信号方向为{np.rad2deg(theta_est)}")
            # 排除错误解算
            if theta_est is not None:
                errors.append(np.abs(theta_true - theta_est))
                angles.append((theta_true, theta_est))
        # 计算平均误差
        mean_error = np.mean(errors) * 180 / np.pi
        print("mean error: {:.2f} degree".format(mean_error))
        return mean_error, angles

if __name__ == "__main__":
    doa = DOA()
    # SNR_list, rmse_list = doa.evaluation_SNR_multi(doa.root_music)
    # SNR_list, rmse_list = doa.evaluation_SNR_multi(doa.esprit)
    # doa.evaluation_SNR_multi((doa.music, doa.root_music, doa.esprit))
    # doa.evaluation_len_s_multi((doa.music, doa.root_music, doa.esprit))
    # doa.plot_SNR_from_file(1000, 40, ("root_music", "esprit"))
    # ("esprit", "music", "root_music")
    # doa.plot_from_file(1000, 15,  ("mean_intensity",), task="len_s")
    # doa.apply_MUSIC()
    #=============
    doa_methods = (doa.music, doa.root_music, doa.esprit, doa.TLS_ESPRIT)
    # doa.SNR_dB = 0
    doa.apply_doa_methods(methods=doa_methods)
    # doa.apply_MUSIC((np.deg2rad(30), np.deg2rad(174)))

    #==================
    # plt.figure()
    # plt.plot(SNR_list, rmse_list)
    # plt.title('RMSE Evaluation for Source Localization')
    # plt.xlabel('SNR/dB')
    # plt.ylabel('RMSE')
    # plt.grid(True)
    # plt.show()
